import torch
import torchaudio
import requests
import matplotlib.pyplot as plt
import os

def test():
    url = "https://pytorch.org/tutorials/_static/img/steam-train-whistle-daniel_simon-converted-from-mp3.wav"
    r = requests.get(url)

    with open('steam-train-whistle-daniel_simon-converted-from-mp3.wav', 'wb') as f:
        f.write(r.content)

    filename = "steam-train-whistle-daniel_simon-converted-from-mp3.wav"
    waveform, sample_rate = torchaudio.load(filename)

    print("Shape of waveform: {}".format(waveform.size()))
    print("Sample rate of waveform: {}".format(sample_rate))

    plt.figure()
    plt.plot(waveform.t().numpy())

    specgram = torchaudio.transforms.Spectrogram()(waveform)

    print("Shape of spectrogram: {}".format(specgram.size()))

    plt.figure()
    plt.imshow(specgram.log2()[0,:,:].numpy())


if __name__ == '__main__':
    # test()
    pwd = os.getcwd()
    test_audio = "dataset\\audio\\collising\\1605472777206Busy1605472777201.wav"# 文件的路径
    test_audio = os.path.join(pwd, test_audio)# 获取整体的路径
    waveform, sample_rate = torchaudio.load(test_audio)
    audio = waveform.t().numpy()
    # print(audio.shape)
    average = audio.mean()
    print(audio.shape)
    pass
    #可视化过程
    plt.figure()
    plt.plot(audio-average)
    plt.show()

    specgram = torchaudio.transforms.Spectrogram()(waveform)
    plt.figure()
    plt.imshow(specgram.log2()[0,:,:].numpy())
    plt.show()